#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import os.path

'''
对已分好类的数据进行按类分割
'''


def main(argv):
    # /Users/David/Desktop/primary_boy_result_6.arff
    # /Users/David/Desktop/boy_cluster
    file_path = argv[1]
    target_path = argv[2]

    print("解析的文件是：", file_path)
    print("目标文件夹是：", target_path)

    common = []  # 存放头部
    data = {}  # 存放数据

    # 读取文件，先存公共部分，再存数据部分
    with open(file_path, 'r') as f:
        line = f.readline()
        while line:
            if line.startswith('@relation') or line.startswith('@attribute') or line.startswith('@data'):
                common.append(line)
            else:
                if line != '\n':
                    fields = line.strip('\n').split(',')
                    key = fields[-1]
                    if key in data:
                        data.get(key).append(line)
                    else:
                        data[key] = [line]
            line = f.readline()

    # 公共部分需要两个换行
    common.insert(1, '\n')
    common.insert(-1, '\n')

    # 开始按类别写入文件
    for key in data.keys():
        file_name = key + '.arff'
        target_file = os.path.join(target_path, file_name)
        print('生成文件：', target_file)
        with open(target_file, 'w') as fp:
            # 写入公共部分
            for line in common:
                # 替换标题
                if line.startswith('@relation'):
                    line = '@relation ' + key + '\n'
                # 替换类别限制
                if line.startswith('@attribute Cluster'):
                    line = '@attribute Cluster {' + key + '}\n'
                fp.write(line)
            # 写入数据部分
            for line in data[key]:
                fp.write(line)


if __name__ == '__main__':
    argv = sys.argv
    main(argv)
